from pathlib import Path

from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict

from ultralytics import YOLO


class Settings(BaseSettings):
    model_config = SettingsConfigDict(env_prefix="yolo_train_")

    weight_path: Path = Field(default=Path("yolov8n.pt"))
    dataset_path: Path = Field(default=Path("dataset.yml"))
    epochs: int = Field(default=100)
    batch_size: int = Field(default=16)
    learning_rate: float = Field(default=0.001)


def main():
    settings = Settings()

    model = YOLO(settings.weight_path)
    model.train(
        data=settings.dataset_path,
        epochs=settings.epochs,
        batch_size=settings.batch_size,
        lr=settings.learning_rate,
    )


if __name__ == "__main__":
    main()
